package com.yitiao.socket.chatroom.nio;

import com.yitiao.socket.chatroom.CommonConst;
import com.yitiao.socket.chatroom.bio.ChatHandler;
import lombok.Data;
import lombok.SneakyThrows;

import java.io.BufferedWriter;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.Buffer;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

@Data
public class NioChatServer {

    // 不再使用 ServerSocket 而是用 Channel 通信。
    private ServerSocketChannel server;
    private Selector selector;
    private ByteBuffer reader;
    private ByteBuffer writer;

    private Charset charset;

    private volatile ConcurrentHashMap<Integer, BufferedWriter> connectedClientMap;

    @SneakyThrows
    public NioChatServer() {
        this.connectedClientMap = new ConcurrentHashMap<>();
        this.reader = ByteBuffer.allocate(CommonConst.BUFFER);
        this.writer = ByteBuffer.allocate(CommonConst.BUFFER);
        this.charset = StandardCharsets.UTF_8;
    }

    public void forwardMessage(SocketChannel client ,String msg) throws IOException {
        for (SelectionKey key: selector.keys()) {
            Channel connectedClient = key.channel();
            if (connectedClient instanceof ServerSocketChannel) {
                continue;
            }

            if (key.isValid() && !client.equals(connectedClient)) {
                writer.clear();
                writer.put(charset.encode("client["+client.socket().getPort()+"]:" + msg));
                writer.flip();
                while (writer.hasRemaining()) {
                    ((SocketChannel)connectedClient).write(writer);
                }
            }
        }
    }

    public synchronized void close() {
        if (server != null) {
            try {
                server.close();
                System.out.println("server has closed.");
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    @SneakyThrows
    public void start(){
        // 创建 server
        server = ServerSocketChannel.open();
        server.configureBlocking(false);  // 关闭阻塞
        server.bind(new InetSocketAddress(CommonConst.PORT));
        // 创建 selector
        selector = Selector.open();
        // 注册到 selector ，监听 accept 事件，即有客户端加入连接
        server.register(selector, SelectionKey.OP_ACCEPT);
        System.out.println("server has started,port is "+server.socket().getLocalPort());

        while (true){
            selector.select();
            Set<SelectionKey> keys = selector.selectedKeys();
            keys.forEach(this::selectionKeyHandler);
            keys.clear();
        }
    }

    @SneakyThrows
    private void selectionKeyHandler(SelectionKey key) {
        if (key.isAcceptable()){
            ServerSocketChannel server = (ServerSocketChannel) key.channel();
            SocketChannel client = server.accept();
            client.configureBlocking(false);
            client.register(selector, SelectionKey.OP_READ);
            System.out.println("client has joined chat["+client.socket().getPort()+"]");
        }
        if (key.isReadable()){
            SocketChannel client = (SocketChannel) key.channel();
            String msg = receive(client);
            if (msg.isEmpty()){
                key.channel();
                selector.wakeup();
            }else {
                System.out.println("client["+client.socket().getPort()+"]:"+msg);

                forwardMessage(client, msg);
                if (msg.equals(CommonConst.QUIT)){
                    key.cancel();
                    selector.wakeup();
                }
            }

        }
    }

    private String receive(SocketChannel client) throws IOException {
        reader.clear();
        while(client.read(reader) > 0);
        reader.flip();
        return String.valueOf(charset.decode(reader));
    }

    public static void main(String[] args) {
        NioChatServer server = new NioChatServer();
        server.start();
    }
}
